#
# Created: 2025-07-28

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
from scipy import stats
from scipy.stats import gaussian_kde

# 设置字体
plt.rcParams['font.family'] = 'Times New Roman'

# 读取 CSV 数据
df = pd.read_csv("random_data.csv")
x = df['x'].values
y = df['y'].values

groups = ['Group 1', 'Group 2']
colors = ['cornflowerblue', 'lightcoral']
fig = plt.figure(figsize=(10, 5))
outer = gridspec.GridSpec(1, 2, wspace=0.1)

main_axes = []

for i, (group, color) in enumerate(zip(groups, colors)):
    width_ratios = [1, 5] if i == 0 else [5, 1]
    inner = gridspec.GridSpecFromSubplotSpec(
        2, 2,
        subplot_spec=outer[i],
        width_ratios=width_ratios,
        height_ratios=[1, 5],
        hspace=0.05, wspace=0.05
    )

    if i == 0:
        ax_histx = plt.Subplot(fig, inner[0, 1])
        ax_histy = plt.Subplot(fig, inner[1, 0])
        ax_main = plt.Subplot(fig, inner[1, 1])
    else:
        ax_histx = plt.Subplot(fig, inner[0, 0])
        ax_histy = plt.Subplot(fig, inner[1, 1])
        ax_main = plt.Subplot(fig, inner[1, 0], sharey=main_axes[0])

    main_axes.append(ax_main)

    # 主图：散点图 + 回归线 + 置信区间
    slope, intercept, r_value, _, _ = stats.linregress(x, y)
    line_x = np.linspace(np.min(x), np.max(x), 100)
    line_y = intercept + slope * line_x
    y_pred = slope * x + intercept
    residual = y - y_pred
    dof = x.size - 2
    t_val = stats.t.ppf(0.975, dof)
    se_line = np.sqrt(np.var(residual) / len(x) + (line_x - np.mean(x)) ** 2 * np.var(x) / ((len(x) - 1) * np.var(x)))
    ci = t_val * se_line

    ax_main.scatter(x, y, s=15, alpha=0.4, color=color)
    ax_main.plot(line_x, line_y, color='black', lw=1.5)
    ax_main.fill_between(line_x, line_y - ci, line_y + ci, color='gray', alpha=0.3)
    eq_text = f'y = {slope:.2f}x + {intercept:.2f}\nR² = {r_value ** 2:.3f}'
    ax_main.text(0.05, 0.95, eq_text, transform=ax_main.transAxes,
                 ha='left', va='top', bbox=dict(facecolor='white', alpha=0.8, edgecolor='none'))

    ax_main.tick_params(which='both', direction='in', top=True, right=True)
    ax_main.tick_params(which='major', length=6)
    ax_main.tick_params(which='minor', length=3)
    ax_main.minorticks_on()

    if i == 1:
        ax_main.set_yticklabels([])
        ax_main.set_ylabel('')
    else:
        ax_main.set_ylabel('y axis', labelpad=20)

    ax_main.set_xlabel('', labelpad=10)
    ax_main.grid(True, linestyle='--', alpha=0.2, which='both')

    # 上方直方图 + KDE
    counts, bins, _ = ax_histx.hist(x, bins=20, color=color, edgecolor='k', alpha=0.6, density=True)
    kde = gaussian_kde(x)
    xx = np.linspace(np.min(x), np.max(x), 200)
    ax_histx.plot(xx, kde(xx), color=color, lw=2)
    ax_histx.set_xticks([])
    ax_histx.set_yticks([])
    ax_histx.set_title(f"{chr(65 + i)}. {group}", loc='left', pad=10)

    # 右侧脊椎图 + KDE
    counts, bins = np.histogram(y, bins=20, density=True)
    centers = (bins[:-1] + bins[1:]) / 2
    ax_histy.barh(y=centers, width=counts, height=(bins[1] - bins[0]),
                  color=color, edgecolor='k', alpha=0.6)
    kde_y = gaussian_kde(y)
    yy = np.linspace(np.min(y), np.max(y), 200)
    kde_val = kde_y(yy)
    ax_histy.plot(kde_val, yy, color=color, lw=2)

    ax_histy.tick_params(axis='y', which='both', direction='in',
                         left=True, right=False, labelleft=True, labelright=False)
    ax_histy.tick_params(which='major', length=6)
    ax_histy.tick_params(which='minor', length=3)
    ax_histy.minorticks_on()

    if i == 0:
        ax_histy.yaxis.tick_left()
        ax_histy.yaxis.set_label_position('left')
        ax_histy.spines['right'].set_visible(True)
        ax_histy.spines['top'].set_visible(True)
        ax_histy.invert_xaxis()
    else:
        ax_histy.yaxis.tick_right()
        ax_histy.yaxis.set_label_position('right')
        ax_histy.spines['left'].set_visible(True)
        ax_histy.spines['top'].set_visible(True)

    ax_histy.set_xticks([])
    ax_histy.grid(True, linestyle='--', alpha=0.2, axis='y')

    fig.add_subplot(ax_histx)
    fig.add_subplot(ax_histy)
    fig.add_subplot(ax_main)

fig.suptitle('Scatter plots with marginal distributions and regression lines', fontsize=14, y=1.02)
save_path = '边际条形图_CSV版本.png'
plt.savefig(save_path, dpi=600, bbox_inches='tight')
plt.show()
print(f"✅ 已生成图，文件保存为：{save_path}")